#include "HdrProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "../HdrSimple/Scene/LocalShader.h"

void HdrProcess::OnInit()
{
	mIsDirty = false;
	mCubemapProcess = "IrradianceCubemapProcess";

	mExposure = 2.5f;
	mGamma = 2.2f;

	float cubeSize = mGameScene->GetCamera().GetFarZ();

	Ubo ubo;
	ubo.proj = mGameScene->GetCamera().GetProj();
	ubo.model = mGameScene->GetCamera().GetView();

	Exposure exposure;
	exposure.exposure = mExposure;
	exposure.gamma = mGamma;

	ENQUEUE_RENDER_COMMAND(BuildModel)([this, cubeSize, ubo, exposure](VkCommandBuffer commandBuffer)
		{
			GeometryGenerator geometryGenerator;
			GeometryGenerator::MeshData cubeData = geometryGenerator.CreateCube(cubeSize, cubeSize, cubeSize, 0);

			std::vector<glm::vec3> vertexs;
			vertexs.resize(cubeData.Vertices.size());
			for (uint32 i = 0; i < cubeData.Vertices.size(); ++i)
			{
				vertexs[i] = cubeData.Vertices[i].Position;
			}

			mVertexBuffer = mVertexBuffer = VertexBuffer::Builder()
				.SetData(vertexs.data())
				.SetSize(vertexs.size() * sizeof(glm::vec3))
				.Build();

			mIndexBuffer = IndexBuffer::Builder()
				.SetData(cubeData.Indices32.data())
				.SetSize(cubeData.Indices32.size() * sizeof(uint32))
				.SetIndexCount(cubeData.Indices32.size())
				.Build();

			Ubo tempUbo = ubo;

			mUboBuffer = CustomBuffer::Builder()
				.SetData(&tempUbo)
				.SetSize(sizeof(Ubo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build(this, "UboBuffer");

			Exposure tempExposure = exposure;

			mExposureBuffer = CustomBuffer::Builder()
				.SetData(&tempExposure)
				.SetSize(sizeof(Exposure))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build(this, "ExposureBuffer");

			mDescriptor = RHIDescriptor::Builder()
				.PushShader<HdrVertexShader>()
				.PushShader<HdrFragmentShader>()
				.SetDescriptorCount(0, 1)
				.Build();

			RefCountPtr<RHITexture> cubemapTexture = GetTexture("IrradianceCubemapProcess", "CubemapTexture");
			RefCountPtr<RHISampler> cubemapSampler = GetSampler("IrradianceCubemapProcess", "CubemapSampler");

			mDescriptor->SetValue("ubo", mUboBuffer);
			mDescriptor->SetValue("samplerEnvMap", cubemapTexture, cubemapSampler);
			mDescriptor->SetValue("exposureUbo", mExposureBuffer);
			mDescriptorMask = mDescriptor->AllocateDescriptor();

			RefCountPtr<RHITexture> colorRenderTarget = GetTexture("LocalProcess", "ColorRenderTarget");
			RefCountPtr<RHITexture> depthRenderTarget = GetTexture("LocalProcess", "DepthRenderTarget");
			RefCountPtr<RHIRenderPass> renderPass = GetRenderPass("LocalProcess", "RenderPass");

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(colorRenderTarget)
				.SetDepthRenderTarget(depthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetCullMode(VK_CULL_MODE_BACK_BIT)
				.SetRenderPass(renderPass->GetRenderPass())
				.Build();
		});
}

void HdrProcess::OnTick(float deltaTime)
{
	Ubo ubo;
	ubo.proj = mGameScene->GetCamera().GetProj();
	ubo.model = mGameScene->GetCamera().GetView();

	Exposure exposure;
	exposure.exposure = mExposure;
	exposure.gamma = mGamma;

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, ubo, exposure](VkCommandBuffer commandBuffer)
		{
			Ubo tempUbo = ubo;
			Exposure tempExposure = exposure;

			mUboBuffer->WriteMemory(&tempUbo);
			mExposureBuffer->WriteMemory(&tempExposure);

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mDescriptorMask.GetDescriptorSet(), 0, nullptr);

			VkBuffer vertexBuffer = mVertexBuffer->GetBuffer();
			VkDeviceSize offsets = 0;
			vkCmdBindVertexBuffers(commandBuffer, 0, 1, &vertexBuffer, &offsets);

			VkBuffer indexBuffer = mIndexBuffer->GetBuffer();
			vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);

			vkCmdDrawIndexed(commandBuffer, mIndexBuffer->GetIndexCount(), 1, 0, 0, 0);

			// vkCmdEndRenderPass(commandBuffer);
		});
}

void HdrProcess::OnResize()
{
}

void HdrProcess::OnFrameEnd()
{
	if (mIsDirty)
	{
		vkDeviceWaitIdle(GVulkanRHI->GetLogicalDevice());

		if (mCubemapProcess == "IrradianceCubemapProcess")
		{
			RefCountPtr<RHITexture> cubemapTexture = GetTexture("IrradianceCubemapProcess", "CubemapTexture");
			RefCountPtr<RHISampler> cubemapSampler = GetSampler("IrradianceCubemapProcess", "CubemapSampler");
			mDescriptor->SetValue("samplerEnvMap", cubemapTexture, cubemapSampler);
		}
		else if (mCubemapProcess == "IrradianceProcess")
		{
			RefCountPtr<RHITexture> cubemapTexture = GetTexture("IrradianceProcess", "IrradianceCube");
			RefCountPtr<RHISampler> cubemapSampler = GetSampler("IrradianceProcess", "IrradianceSampler");
			mDescriptor->SetValue("samplerEnvMap", cubemapTexture, cubemapSampler);
		}
		else if (mCubemapProcess == "PrefilterenvmapProcess")
		{
			RefCountPtr<RHITexture> cubemapTexture = GetTexture("PrefilterenvmapProcess", "PrefilterenvmapCube");
			RefCountPtr<RHISampler> cubemapSampler = GetSampler("PrefilterenvmapProcess", "PrefilterenvmapSampler");
			mDescriptor->SetValue("samplerEnvMap", cubemapTexture, cubemapSampler);
		}
		mDescriptor->UpdateDescriptor(mDescriptorMask);

		mIsDirty = false;
	}
}

void HdrProcess::UpdateCubemap(const std::string& processName)
{
	if (mCubemapProcess == processName)
		return;

	mCubemapProcess = processName;
	mIsDirty = true;
}
